# -*- coding: utf-8 -*-
"""
Created on Tue Mar 18 18:43:10 2025

@author: lenovo
"""

import pandas as pd
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
import seaborn as sns
from global_land_mask import globe
from gsj_typhoon import tydat,see


# 不同起报时间文件路径
dir_dsr=r"D:\met_data\dusurui"
name_date=os.listdir(dir_dsr)
dir_date= [os.path.join(dir_dsr,f) for f in name_date ]
RIstd = 7  # std = 15   
member_num = 52  # 0-51号




#用于记录所有数据；所有起报时间内的所有成员以及对应成员ID
num = [[None for _ in range(member_num)] for _ in range(len(name_date))]
for i in range(len(dir_date)):
    date = dir_date[i]
    # 得到某一起报时刻的所有集合成员数据
    name_all = [f for f in os.listdir(date) if f.startswith("TRACK")]
    sorted_names = sorted(name_all, key=lambda x: int(x.split('TRACK_ID_')[-1]))
    path_all = [os.path.join(date, f) for f in sorted_names ]
    #通过列表记录所有成员，借助类来封装处理方式
    for path in path_all:
        t=tydat(path,RIstd)
        nn = int(path.split("TRACK_ID_")[-1] )
        # 记录单个成员数据以及成员序号
        num[i][nn] = t



###  创建绘图数组   ###
a=[];b=[]
for s in range(0,len(name_date)):
    start_time=datetime.strptime(name_date[s],"%Y%m%d%H")
    for i in range(0,len(num[s])):
        if num[s][i] !=None :
            a.append( max(num[s][i].time ))
            b.append( max(num[s][i].time) - start_time )
max_enddate = max(a)
max_dayrange = max(b)
draw_data = np.zeros( (len(name_date),max_dayrange.days+1)  , dtype=int )



### 填充draw_data ###
tt,dd = draw_data.shape



#集合成员
for i in range(member_num):
    #起报
    for t in range(tt):
        #预报时效
        for d in range(dd):
            start_time = datetime.strptime(name_date[t],"%Y%m%d%H")
            cal_time = start_time + d*timedelta(days=1)
            calend_time = cal_time + timedelta(hours=18)
            #成员
            if num[t][i]!=None:
                data = num[t][i]
                RI = data.num_rapidgrow()
                dayRI = RI[ (data.time>=cal_time) & (data.time<=calend_time) ]
                n=0
                if np.sum(dayRI>=1) :
                    n+=1
                draw_data[t][d]=n
            
    # draw
    plt.figure(figsize=(8, 6))
    ax = sns.heatmap(draw_data.T[:15],annot=True,annot_kws={'size': 6},fmt="d", cmap="YlGnBu", cbar_kws={'label': 'number of rapid intensification'}, linewidths=0.5)

    ax.set_ylabel("Forecast realative to start time", fontsize=12)
    ax.set_xlabel("Start time", fontsize=12)

    # y坐标作为时间轴，先处理好datetime对象
    # plt.tick_params(axis='both',which='both',labelsize=10)
    ax.set_xticklabels(name_date, rotation=45) 
    plt.yticks(rotation=0)
    ylabels=[str(i)+"d"  for i in range(1,16)]
    ax.set_yticklabels(ylabels)
    ax.set_title(dir_dsr.split("\\")[-1]+f"-RI{RIstd} ID{i}" )
    # plt.savefig("./heatmap_days_startime",dpi=900)
    plt.show()
            
        
        
    




